package main

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/hyperledger/fabric/bccsp"
	"github.com/hyperledger/fabric/bccsp/factory"
	"github.com/hyperledger/fabric/core/chaincode/shim"
	"github.com/hyperledger/fabric/core/chaincode/shim/ext/cid"
	"github.com/hyperledger/fabric/core/chaincode/shim/ext/entities"
	"github.com/hyperledger/fabric/protos/peer"
)

type SimpleAsset struct {
	bccspInst bccsp.BCCSP
}

func main() {

	factory.InitFactories(nil)
	cc := SimpleAsset{bccspInst:factory.GetDefault()}
	if err := shim.Start(&cc); err != nil {
		fmt.Println("fail to start chaincode:", err)
	}
}

func (s *SimpleAsset) Init(stub shim.ChaincodeStubInterface) peer.Response {
	return shim.Success(nil)
}

func (s *SimpleAsset) Invoke(stub shim.ChaincodeStubInterface) peer.Response {
	fcn, args := stub.GetFunctionAndParameters()
	switch fcn {
	case "set":
		return setValue(stub, args)
	case "setPrivate":
		return setPrivateValue(stub, args)
	case "get":
		return getValue(stub, args)
	case "getPrivate":
		return getPrivateValue(stub, args)
	case "delete":
		return deleteValue(stub, args)
	case "test":
		return testApi(stub, args)
	case "setEnc":
		return s.setValueEncrypt(stub, args)
	case "getDec":
		return s.getValueDecrypty(stub, args)
	case "setEncSign":
		return s.setEncryptSign(stub, args)
	case "decVerify":
		return s.decryptVerify(stub, args)
	}
	return shim.Error("undefined function " + fcn)
}

func testApi(stub shim.ChaincodeStubInterface, args []string) peer.Response {
	clientID, err := cid.New(stub)
	if err != nil {
		return shim.Error(err.Error())
	}

	mspid, _ := clientID.GetMSPID()
	fmt.Println("mspid: ", mspid)

	val, found, err := clientID.GetAttributeValue("role")
	if err != nil {
		return shim.Error(err.Error())
	}
	fmt.Println("role:", found, val)

	val, found, _ = clientID.GetAttributeValue("qq")
	fmt.Println("qq:", val, found)
	val, found, _ = clientID.GetAttributeValue("phoneNo")
	fmt.Println("phoneNO", val, found)

	tmap, err := stub.GetTransient()
	if err != nil {
		return shim.Error(err.Error())
	}
	fmt.Println("transient: ", tmap)
	return shim.Success(nil)
}

func setValue(stub shim.ChaincodeStubInterface, args []string) peer.Response {
	if len(args) != 2 {
		return shim.Error(fmt.Sprintf("argument invalid, need %d, actual id %d", 2, len(args)))
	}
	fmt.Println("set value: ", args[0])
	value, err := stub.GetState(args[0])
	if err != nil {
		return shim.Error("fail to get state of " + args[0] + " with error:" + err.Error())
	} else if value != nil {
		return shim.Error("key " + args[0] + " exist")
	}
	stub.PutState(args[0], []byte(args[1]))
	return shim.Success(nil)
}

// need fabric 1.2 above
func setPrivateValue(stub shim.ChaincodeStubInterface, args []string) peer.Response {
	if len(args) != 3 {
		return shim.Error(fmt.Sprintf("argument invalid, need %d, actual id %d", 2, len(args)))
	}
	collName := args[0]
	fmt.Println("set collection name:", collName)
	value, err := stub.GetPrivateData(collName, args[1])
	if err != nil {
		return shim.Error("fail to get state of " + args[0] + " with error:" + err.Error())
	} else if value != nil {
		return shim.Error("key " + args[1] + "exist")
	}
	stub.PutPrivateData(collName, args[1], []byte(args[2]))
	return shim.Success(nil)
}

func getValue(stub shim.ChaincodeStubInterface, args []string) peer.Response {
	response := new(bytes.Buffer)
	bArrayMemberAlreadyWritten := false
	response.WriteString("[")
	for _, a := range args {
		val, err := stub.GetState(a)
		if err != nil {
			return shim.Error("key " + a + " does not exist")
		}
		if bArrayMemberAlreadyWritten == true {
			response.WriteString(",")
		}
		response.WriteString(a + ":" + string(val))
		bArrayMemberAlreadyWritten = true
	}
	response.WriteString("]")
	return shim.Success(response.Bytes())
}

// need fabric 1.2 above
func getPrivateValue(stub shim.ChaincodeStubInterface, args []string) peer.Response {
	if len(args) < 1 {
		return shim.Error("argument invalid, need at least 1")
	}
	response := new(bytes.Buffer)
	bArrayMemberAlreadyWritten := false
	collName := args[0]
	fmt.Println("get collection name:", collName)
	response.WriteString("[")
	for _, a := range args[1:] {
		val, err := stub.GetPrivateData(collName, a)
		if err != nil {
			//return shim.Error("key " + a + " does not exist")
			fmt.Println("key " + a + " does not exist")
			continue
		}
		if bArrayMemberAlreadyWritten == true {
			response.WriteString(", ")
		}
		response.WriteString(a + ":" + string(val))
		bArrayMemberAlreadyWritten = true
	}
	response.WriteString("]")
	return shim.Success(response.Bytes())
}

func deleteValue(stub shim.ChaincodeStubInterface, args []string) peer.Response {
	if len(args) < 1 {
		return shim.Error("argument invalid, need at least 1")
	}

	var resp []string
	for _, a := range args[1:] {
		err := stub.DelState(a)
		if err != nil {
			resp = append(resp, "delete key "+a+" fail with error:"+err.Error())
		}
	}
	if len(resp) == 0 {
		return shim.Success([]byte("delete ok"))
	}
	v, _ := json.Marshal(resp)
	return shim.Success(v)
}


// set value that be encrypt
// ENCKEY=`openssl rand 32 | base64`
// peer chaincode invoke -n enccc -C mychannel -c '{"Args":["setEnc","key1","value1"]}' --transient "{\"encKey\":\"$ENCKEY\"}"
// IV=`openssl rand 16 | base64`
// peer chaincode invoke -n enccc -C mychannel -c '{"Args":["setEnc","key2","value2"]}' --transient "{\"encKey\":\"$ENCKEY\",\"iv\":\"$IV\"}"
func (sa SimpleAsset) setValueEncrypt (stub shim.ChaincodeStubInterface, args []string) peer.Response {
	key := args[0]
	//encryptKey := args[1]
	value := args[1]
	tmap, err := stub.GetTransient()
	if err != nil {
		return shim.Error(err.Error())
	}
	encKey, _ := tmap["encKey"]
	if encKey == nil {
		return shim.Error("encKey not found")
	}
	encEntry, err := entities.NewAES256EncrypterEntity("ID", sa.bccspInst, encKey, tmap["iv"])
	if err != nil {
		return shim.Error(err.Error())
	}
	cipherText, err := encEntry.Encrypt([]byte(value))
	if err != nil {
		return shim.Error("encrypt:" + err.Error())
	}
	stub.PutState(key, cipherText)
	// TODO
	return shim.Success(nil)
}

// DECKEY=$ENCKEY
// peer chaincode query -n enccc -C mychannel -c '{"Args":["getDec","key1"]}' --transient "{\"decKey\":\"$DECKEY\"}"
func (sa SimpleAsset) getValueDecrypty(stub shim.ChaincodeStubInterface, args []string) peer.Response {
	key := args[0]
	tmap, err := stub.GetTransient()
	if err != nil {
		return shim.Error(err.Error())
	}
	decKey, _ := tmap["denKey"]
	if decKey == nil {
		return shim.Error("decKey not found")
	}
	decEntry, err := entities.NewAES256EncrypterEntity("ID", sa.bccspInst, decKey, tmap["iv"])
	if err != nil {
		return shim.Error("new encrypt entry error:" + err.Error())
	}
	ciphertext, _ := stub.GetState(key)
	if len(ciphertext) == 0 {
		return shim.Error("no ciphertext to decrypt")
	}

	decData, err := decEntry.Decrypt(ciphertext)
	if err != nil {
		return shim.Error("decrypt error:" + err.Error())
	}
	return shim.Success(decData)
}


func encryptAndPutState(stub shim.ChaincodeStubInterface, ent entities.Encrypter, key string, value []byte) error {
	// at first we use the supplied entity to encrypt the value
	ciphertext, err := ent.Encrypt(value)
	if err != nil {
		return err
	}

	return stub.PutState(key, ciphertext)
}

func getStateAndDecrypt(stub shim.ChaincodeStubInterface, ent entities.Encrypter, key string) ([]byte, error) {
	ciphertext, err := stub.GetState(key)
	if err != nil {
		return nil, err
	}

	if len(ciphertext) == 0 {
		return nil, errors.New("no ciphertext to decrypt")
	}

	return ent.Decrypt(ciphertext)
}

// SIGKEY=`openssl ecparam -name prime256v1 -genkey | tail -n5 | base64 -w0`
// peer chaincode invoke -n enccc -C mychannel -c '{"Args":["setEncSign","key3","value3"]}' --transient "{\"encKey\":\"$ENCKEY\",\"sig\":\"$SIGKEY\"}"
func (sa SimpleAsset) setEncryptSign(stub shim.ChaincodeStubInterface, args []string) peer.Response {
	tMap, err := stub.GetTransient()
	if err != nil {
		return shim.Error(err.Error())
	}

	encKey, _ := tMap["encKey"]
	if encKey == nil {
		return shim.Error("not found encKey")
	}
	sig, _ := tMap["sig"]
	if sig == nil {
		return shim.Error("not found verify key")
	}
	encSignerEntry, err := entities.NewAES256EncrypterECDSASignerEntity("ID", sa.bccspInst, encKey, sig)
	if err != nil {
		return shim.Error(fmt.Sprintf("entities.NewAES256EncrypterEntity failed, err %s", err))
	}
	key := args[0]
	cleartextValue := []byte(args[1])

	msg := &entities.SignedMessage{Payload: cleartextValue, ID: []byte(encSignerEntry.ID())}
	err = msg.Sign(encSignerEntry)
	if err != nil {
		return shim.Error(err.Error())
	}
	signedBytes, err := msg.ToBytes()
	if err != nil {
		return shim.Error(err.Error())
	}
	err = encryptAndPutState(stub, encSignerEntry, key, signedBytes)
	if err != nil {
		return shim.Error(err.Error())
	}
	return shim.Success([]byte("success"))
}

// VERKEY=$SIGKEY
// peer chaincode query -n enccc -C mychannel -c '{"Args":["decVerify","key3"]}' --transient "{\"deckey\":\"$DECKEY\",\"verify\":\"$VERKEY\"}"
func (sa SimpleAsset) decryptVerify(stub shim.ChaincodeStubInterface, args []string) peer.Response {
	tMap, err := stub.GetTransient()
	if err != nil {
		return shim.Error(err.Error())
	}
	decKey, _ := tMap["decKey"]
	if decKey == nil {
		return shim.Error("not found encKey")
	}
	verifyKey, _ := tMap["verify"]
	if verifyKey == nil {
		return shim.Error("not found verify key")
	}
	ent, err := entities.NewAES256EncrypterECDSASignerEntity("ID", sa.bccspInst, decKey, verifyKey)
	if err != nil {
		return shim.Error(fmt.Sprintf("entities.NewAES256DecrypterEntity failed, err %s", err))
	}
	key := args[0]
	// get value and verify
	val, err := getStateAndDecrypt(stub, ent, key)
	if err != nil {
		return shim.Error(err.Error())
	}
	msg := &entities.SignedMessage{}
	err = msg.FromBytes(val)
	if err != nil {
		return shim.Error(err.Error())
	}
	// we verify the signature
	ok, err := msg.Verify(ent)
	if err != nil {
		return shim.Error(err.Error())
	} else if !ok {
		return shim.Error("invalid signature")
	}
	return shim.Success(msg.Payload)
}